package probability;

import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;
import java.util.Random;

import basics.DenseMatrix;
import basics.DenseVector;
import basics.Matrix;
import basics.Vector;
import basics.VectorMatrix;
import basics.VectorMatrixUtils;

public class GaussianDistribution extends ProbabilityDistribution {

	// private Vector _mean;
	// private DenseMatrix _var;

	// TODO: replace all the the params references to the mu and sigmas

	private static final int SAMPLE_SIZE = 1000;

	public GaussianDistribution(ArrayList<VectorMatrix> params) {
		super(params);
	}

	public GaussianDistribution() {
		super();
		_params.add(new DenseVector(new double[] { 0. }));
		_params.add(new DenseMatrix(new double[][] { { 1. } }));
	}

	public void setParam(int i, double d) {
		VectorMatrix vm;
		if (i == 0) {
			vm = new DenseVector(new double[] { d });
		} else {
			vm = new DenseMatrix(new double[][] { { d } });
		}
		setParam(i, vm);
	}

	public GaussianDistribution(double mean, double var) {
		super();
		// _mean = new DenseVector(1);
		// _mean.set(mean, 0);
		// _params.add(_mean);
		// _var = new DenseMatrix(1, 1);
		// _var.set(var, 0, 0);
		// _params.add(_var);

		_params.add(new DenseVector(new double[] { mean }));
		_params.add(new DenseMatrix(new double[][] { { var } }));
	}

	@Override
	public double measure(Vector x) {
		Vector mean = (Vector) getParam(0);
		Matrix var = (Matrix) getParam(1);
		Matrix t1 = (Matrix) var.times(var).times(2);
		Vector t2 = (Vector) x.minus(mean);
		double p = 1. / Math.sqrt((Double) t1.times(Math.PI).get(0, 0));
		p = p * Math.exp(-1 * ((Double) t1.timesEqual(t1).get(0, 0)) / (Double) t2.get(0));

		return p;
	}

	public double getMu() {
		return getParam(0).get(0);
	}

	public double getSigma2() {
		return getParam(1).get(0, 0);
	}

	public void setMu(double mu) {
		getParam(0).set(mu, 0);
	}

	public void setSigma2(double sigma) {
		getParam(1).set(Math.max(sigma, 0.1), 0, 0);
	}

	public double getRawMoment(int order) {
		switch (order) {
		case 1:
			return getMu();
		case 2:
			return Math.pow(getMu(), 2) + getSigma2();
		case 3:
			return Math.pow(getMu(), 3) + 3 * getMu() * getSigma2();
		default:
			new Exception("This Moment is not defined till now").printStackTrace();
		}

		return -1;
	}

	public void setPiTaw(double pi, double taw) {
		setSigma2(1. / pi);
		setMu(taw / pi);
	}

	public double getTaw() {
		return getPi() * getMu();
	}

	public double getPi() {
		return 1. / getSigma2();
	}

	public double Q(Vector x) {
		Vector mean = (Vector) getParam(0);
		Matrix var = (Matrix) getParam(1);
		return Q(x.get(0), mean.get(0), var.get(0, 0));
	}

	public static double erf(double x) {
		double a = 0.140012;
		return Math.signum(x) * Math.sqrt(1 - Math.exp(-x * x * ((4. / Math.PI + a * x * x) / (1. + a * x * x))));
	}

	public static double Q(double x, double mu, double sigma) {
		return .5 * (1 + erf((x - mu) / sigma * Math.sqrt(2)));
	}

	public static GaussianDistribution product(GaussianDistribution... distributions) {
		return product(Arrays.asList(distributions));
	}

	public DenseVector sample(int sampleSize) {
		DenseVector d = new DenseVector(sampleSize);
		Random r = new Random();
		double s = Math.sqrt(getSigma2());
		for (int i = 0; i < sampleSize; i++) {
			d.set(getMu() + s * r.nextGaussian(), i);
		}

		return d;
	}

	// TODO: Correct all the products
	public static GaussianDistribution product(List<GaussianDistribution> distributions) {
		GaussianDistribution gd = null;
		for (GaussianDistribution g : distributions) {
			if (gd == null) {
				gd = new GaussianDistribution(g.getMu(), g.getSigma2());
			} else {
				gd.setMu((gd.getMu() * g.getSigma2() + g.getMu() * gd.getSigma2()) / (gd.getSigma2() + g.getSigma2()));
				gd.setSigma2((gd.getSigma2() * g.getSigma2()) / (gd.getSigma2() + g.getSigma2()));
			}
		}

		return gd;
	}

	// This function calculates the product of two gaussian with respect to an indicator function using sampling with
	// specified number of iterations
	public static GaussianDistribution sampleIndicatorProduct(int numberOftimes, GaussianDistribution... gps) {
		GaussianDistribution ret = null;
		for (int i = 0; i < numberOftimes; i++) {
			GaussianDistribution g = sampleIndicatorProduct(gps);
			// System.out.println(g);
			if (ret == null) {
				ret = g;
			} else {
				double s = (double) i / (double) (i + 1.);
				ret.setMu(ret.getMu() * s + g.getMu() * (1. - s));
				ret.setSigma2(ret.getSigma2() * s + g.getSigma2() * (1. - s));
			}
		}
		return ret;
	}

	// This function calculates the product of two gaussian with respect to an indicator function using sampling
	public static GaussianDistribution sampleIndicatorProduct(GaussianDistribution... gps) {
		DenseVector samples = null;

		for (GaussianDistribution gp : gps) {
			DenseVector s = gp.sample(SAMPLE_SIZE);
			if (samples == null) {
				samples = s;
			} else {
				samples.entrywiseProduct(s);
			}
		}

		double mean = VectorMatrixUtils.mean(samples);
		return new GaussianDistribution(mean, Math.pow(VectorMatrixUtils.stddev(samples, mean), 2));
	}

	public static GaussianDistribution sampleIndicatorDivide(int numberOftimes, GaussianDistribution numerator,
			GaussianDistribution denominator) {

		GaussianDistribution ret = null;
		for (int i = 0; i < numberOftimes; i++) {
			DenseVector num = numerator.sample(SAMPLE_SIZE);
			DenseVector den = denominator.sample(SAMPLE_SIZE);
			DenseVector t = (DenseVector) num.divideCopy(den);
			double mean = VectorMatrixUtils.mean(t);
			GaussianDistribution g = new GaussianDistribution(mean, Math.pow(VectorMatrixUtils.stddev(t, mean), 2));
			if (ret == null) {
				ret = g;
			} else {
				double s = (double) i / (double) (i + 1.);
				ret.setMu(ret.getMu() * s + g.getMu() * (1. - s));
				ret.setSigma2(ret.getSigma2() * s + g.getSigma2() * (1. - s));
			}
		}
		return ret;
	}

	public double measure(double d) {
		return measure(new DenseVector(new double[] { d }));
	}

	@Override
	protected ProbabilityDistribution newInstance() {
		return new GaussianDistribution();
	}

	public MultivariateGaussianDistribution getMultivariate(int n) {
		return new MultivariateGaussianDistribution(DenseVector.getValuedVector(n, 1.).times(getMu()),
				(Matrix) DenseMatrix.getIdentityMatrix(n).times(getSigma2()));
	}

	@Override
	public ProbabilityDistribution times(ProbabilityDistribution distribution) {
		return product(this, (GaussianDistribution) distribution);
	}

	public GaussianDistribution divide(ProbabilityDistribution distribution) {
		GaussianDistribution gd = new GaussianDistribution();
		GaussianDistribution g = (GaussianDistribution) distribution;

		double denom = 1. / (g.getSigma2() - this.getSigma2());
		denom = Double.isInfinite(denom) ? 10000 : denom;
		gd.setMu((this.getMu() - g.getMu()) * denom);
		gd.setSigma2((this.getSigma2() * g.getSigma2()) * denom);

		return gd;
	}

	public GaussianDistribution minus(GaussianDistribution g) {
		return add(g.times(-1));
	}

	public GaussianDistribution add(GaussianDistribution g) {
		return new GaussianDistribution(getMu() + g.getMu(), getSigma2() + g.getSigma2());
	}

	public GaussianDistribution times(double coef) {
		return new GaussianDistribution(getMu() * coef, getSigma2() * coef * coef);
	}

	public static GaussianDistribution mean(List<GaussianDistribution> samples, int ignore_count) {
		GaussianDistribution m = new GaussianDistribution(0, 0);
		for (int i = ignore_count; i < samples.size(); i++) {
			GaussianDistribution g = samples.get(i);

			m.setMu(m.getMu() + g.getMu());
			m.setSigma2(m.getSigma2() + g.getSigma2());
		}
		double count = samples.size() - (ignore_count);
		m.setMu(m.getMu() / count);
		m.setSigma2(m.getSigma2() / count);

		return m;
	}

	public static GaussianDistribution sum(GaussianDistribution[] a) {
		GaussianDistribution g = a[0];
		for (int i = 1; i < a.length; i++) {
			g.add(a[i]);
		}
		return g;
	}

}
